package graceful_shutdown

import (
	context2 "context"
	"fmt"
	"log"
	"os"
	"os/signal"
	"server/framework/web/built-in-server/graceful_shutdown/shutdown_signal"
	"server/framework/web/context"
	"server/framework/web/middleware"
	"server/framework/web/server"
	"sync"
	"sync/atomic"
	"time"
)

/*
ShutdownServer 优雅退出
1. 拒绝新的请求
2. 处理完成已有请求
3. 关闭服务
4. 释放资源
5. 超时强制关闭
*/
type ShutdownServer struct {
	reqCnt      int64
	close       int64
	finishReqCh chan struct{}
}

func NewShutdownServer() *ShutdownServer {
	return &ShutdownServer{
		finishReqCh: make(chan struct{}, 1),
	}
}

func (s *ShutdownServer) RequestMiddleware(next middleware.HandleFunc) middleware.HandleFunc {
	return func(c *context.Context) {
		if atomic.LoadInt64(&s.close) > 0 {
			c.RespSystemErr("server already closed")
			return
		}
		atomic.AddInt64(&s.reqCnt, 1)
		next(c)
		atomic.AddInt64(&s.reqCnt, -1)

		cl := atomic.LoadInt64(&s.close)
		reqCnt := atomic.LoadInt64(&s.reqCnt)
		if cl > 0 {
			fmt.Println("当前剩余 reqCnt:", reqCnt)
		}

		if cl > 0 && reqCnt == 0 {
			s.finishReqCh <- struct{}{}
		}
	}
}

func (s *ShutdownServer) RejectRequestBuilder() Hook {
	return func(c context2.Context) {
		atomic.AddInt64(&s.close, 1)
		reqCnt := atomic.LoadInt64(&s.reqCnt)
		if reqCnt == 0 {
			// 无未处理完成请求
			return
		}

		select {
		case <-s.finishReqCh:
			fmt.Println("request finished")
		case <-c.Done():
			fmt.Println("request handle timeout")
		}
	}
}

func (s *ShutdownServer) ServerBuilder(servers ...server.Server) Hook {
	return func(c context2.Context) {
		doneCh := make(chan bool, 1)
		group := sync.WaitGroup{}
		group.Add(len(servers))
		for _, sv := range servers {
			go func(s server.Server) {
				err := s.Shutdown(nil)
				group.Done()
				if err != nil {
					log.Printf("server shutdown error: %v", err)
				}
			}(sv)
		}
		go func() {
			group.Wait()
			doneCh <- true
		}()

		select {
		case <-doneCh:
			fmt.Println("servers shutdown finished")
		case <-c.Done():
			fmt.Println("servers shutdown timeout")
		}
	}
}

type Hook func(c context2.Context)

func WaitForShutdown(hooks ...Hook) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, shutdown_signal.ShutdownSignals...)

	sig := <-signals
	fmt.Printf("Received shutdown signal: %s\n", sig)
	ctx, cancel := context2.WithTimeout(context2.Background(), time.Second*15)

	time.AfterFunc(time.Second*30, func() {
		// 超时强制退出
		fmt.Println("强制退出")
		cancel()
		os.Exit(-1)
	})

	group := sync.WaitGroup{}
	group.Add(len(hooks))
	for _, hk := range hooks {
		go func(h Hook) {
			h(ctx)
			defer group.Done()
		}(hk)
	}

	group.Wait()
	cancel()
	os.Exit(0)
}
